import glob
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
from sklearn.cluster import KMeans
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
from sklearn.manifold import TSNE


def make_corpus() -> list:
    corpus = []
    xlsx_files = glob.glob('../../resource/*.xlsx')
    for xlsx_file in xlsx_files:
        if xlsx_file == '../../resource\\整体词频.xlsx':
            continue
        df = pd.read_excel(xlsx_file)
        data = df['关键词(TF-IDF)']
        list = data.values.tolist()
        # print(list)
        origin_str = ""
        for i in range(len(list)):
            if isinstance(list[i], float) and len(origin_str) != 0:
                origin_str = origin_str[:-1]
                corpus.append(origin_str)
                origin_str = ""
            elif isinstance(list[i], str):
                origin_str = origin_str + list[i] + " "
            else:
                pass
    # print(corpus)
    return corpus


def Kmeans(corpus):
    vectorizer = CountVectorizer()  # 该类会将文本中的词语转换为词频矩阵，矩阵元素a[i][j] 表示j词在i类文本下的词频
    transformer = TfidfTransformer()  # 该类会统计每个词语的tf-idf权值
    tfidf = transformer.fit_transform(vectorizer.fit_transform(corpus))  # 第一个fit_transform是计算tf-idf，第二个fit_transform是将文本转为词频矩阵
    word = vectorizer.get_feature_names_out()  # 获取词袋模型中的所有词语
    weight = tfidf.toarray()  # 将tf-idf矩阵抽取出来，元素a[i][j]表示j词在i类文本中的tf-idf权重

    kmeans = KMeans(n_clusters=5).fit(weight)  # k值可以自己设置，不一定是五类
    centroid_list = kmeans.cluster_centers_  # 聚类中心
    labels = kmeans.labels_  # 聚类标签
    n_clusters_ = len(centroid_list)

    max_centroid = 0
    max_cluster_id = 0
    cluster_menmbers_list = []
    for i in range(0, n_clusters_):
        menmbers_list = []
        for j in range(0, len(labels)):
            if labels[j] == i:
                menmbers_list.append(j)
        cluster_menmbers_list.append(menmbers_list)
    # 聚类结果
    for i in range(0, len(cluster_menmbers_list)):
        print('第' + str(i) + '类' + '---------------------')
        for j in range(0, len(cluster_menmbers_list[i])):
            a = cluster_menmbers_list[i][j]
            print(corpus[a])
    # 散点图数据准备，TSNE降维数，weight为tf-idf矩阵weight = tfidf.toarray()，
    tsne = TSNE(n_components=2)
    decomposition_data = tsne.fit_transform(weight)
    x = []
    y = []
    for i in decomposition_data:
        x.append(i[0])
        y.append(i[1])
    fig = plt.figure(figsize=(10, 10))
    ax = plt.axes()
    plt.scatter(x, y, c=kmeans.labels_, marker=".")
    plt.xticks(())
    plt.yticks(())
    plt.show()
    # plt.savefig('./sample.png', aspect=1)


corpus = make_corpus()
Kmeans(corpus)